from enthought.tvtk.api import tvtk
from enthought.pyface.tvtk import actors
from enthought.tvtk.tools import ivtk
import sys
import os
import itk


if sys.platform == 'linux2':
    fileNameVol = '/home/bryan/Data/vert1.vtk'  #Linux
    fileNameImg = '/home/bryan/jointtrack/branches/bryan/output.png'
else:
    if os.path.exists('C:/Lib/jointtrack/branches/TestData/vert1.vtk'):
        fileNameVol = 'C:/Lib/jointtrack/branches/TestData/vert1.vtk' #Vista
        fileNameImg = 'C:/Lib/jointtrack/branches/bryan/output.png'
    else:
        fileNameVol = 'C:/Data/2D3D/vert1/CT/vert1.vtk' #Windows
        fileNameImg = 'C:/Lib/jointtrack/branches/bryan/output.png' 

class myActors():
        
    def imageActor(self, fn=fileNameImg):
        print "Opening image file: ", fn
        imageReader = tvtk.PNGReader(file_name=fn)
        
        backgroundMapper = tvtk.ImageMapper(input=imageReader.output)
        backgroundMapper.color_window = 256
        backgroundMapper.color_level = 127.5
        imageActor = tvtk.Actor2D()
        imageActor.mapper = backgroundMapper
        
        backgroundPropProperties = imageActor.property
        backgroundPropProperties.display_location = "background"
        
        return imageActor

    def earthActor(self, size=250):
        tmp = actors.earth_actor(radius=size, opacity=1.0)
        return tmp

    def coneActor(self, size=250):
        tmp = actors.cone_actor(height=2*size,radius=size)
        return tmp

    def cubeActor(self, size=250):
        tmp = actors.cube_actor()
        return tmp

    def cylinderActor(self, size=250):
        tmp = actors.cylinder_actor()
        tmp.scale=(250,250,250)
        return tmp

    def arrowActor(self, size=250):
        tmp = actors.arrow_actor()
        return tmp

    def sphereActor(self, size=250):
        tmp = actors.sphere_actor(radius=size)
        return tmp
    
    def volumeActor(self, fn=fileNameVol):
        print "opening volume file: ", fn        
        volume = tvtk.Volume()
        volumeReader = tvtk.StructuredPointsReader()
        volumeReader.file_name = fn
        
        #Default settings for the opacity transfer funciton,
        opacityTransferFunction = tvtk.PiecewiseFunction()        
        opacityTransferFunction.add_point(0.0, 0.0)
        opacityTransferFunction.add_point(70, 0.0)
        opacityTransferFunction.add_point(255, 0.5)
        
        # Create transfer mapping scalar value to color
        colorTransferFunction = tvtk.ColorTransferFunction()
        colorTransferFunction.add_rgb_point(0.0, 0.0, 0.0, 0.0)
        colorTransferFunction.add_rgb_point(72, 1.0, 1.0, 1.0)
        colorTransferFunction.add_rgb_point(255, 1.0, 1.0, 1.0)
        
        # The property describes how the data will look
        volumeProperty = tvtk.VolumeProperty()
        # Setup default properties
        volumeProperty.set_color(colorTransferFunction)
        volumeProperty.set_scalar_opacity(opacityTransferFunction)
        volumeProperty.shade = 1
        volumeProperty.interpolation_type = 'linear'
        
        volumeMapper = tvtk.VolumeTextureMapper3D()
        volumeMapper.set_input(volumeReader.get_output())
        
        volume._set_property(volumeProperty)
        volume._set_mapper(volumeMapper)
        
        volume.origin = volume.center
        
        return volume
    

class util():
    def __init__(self):
        self.volume = []
        self.viewer = []
        self.image = []
        self.style = tvtk.InteractorStyleTrackballActor()
        
    def createViewer(self, mysize=(528,591)):
        """A size of (528,591) for the viewer will create a render window
        of size (512,512)."""
        self.viewer = ivtk.IVTK(size=mysize)
        print "Viewer created, ", mysize

    def addActors(self, actors):
        #The viewer has to be open to access the scene methods
        #TODO: Check to make sure that the viewer has been set
        self.viewer.open()
        #can add a single actor, or pass a tuple of multiple actors
        self.viewer.scene.add_actors(actors)
        style = tvtk.InteractorStyleTrackballActor()
        viewer.scene.interactor.interactor_style = style
        viewer.scene.add_actors(image)
        
    def addVolume(self):
        #The volume and viewer must be set before calling this function
        #TODO: Check to make sure that the viewer has been set 
        print "Adding Volume"
        #print self.viewer.       
        self.viewer.open()
        self.viewer.scene.renderer.add_volume(self.volume)
    
    def addImage(self):
        print "Adding Image"
        self.viewer.open()
        #This is the default style, this probably does not need to be here 
        # unless we decide to use a differnt syle for images only?
        style=tvtk.InteractorStyleTrackballActor()
        self.setStyle(style)
        self.viewer.scene.add_actors(self.image)
##        self.viewer.scene.renderer.reset_camera_clipping_range()
        
    def removeActors(self, actors):
        self.viewer.scene.remove_actors(actors)

    def getAllActors(self):
        """This functon just returns a lst of the built in actors, 
        Not the actors currently added to the renderer!"""
        ACTORS = []
        actorFuncs = [getattr(myActors, x) for x in dir(myActors)
            if x.endswith('Actor')]
        for i in range(len(actorFuncs)):
            # Create the actor.
            func = actorFuncs[i]
            ACTORS.append(func())
        return ACTORS
    
    def doitall(self):
        """"A lazy way to set up a viewer with the default spine volume"""
        m = myActors()
        self.volume = m.volumeActor()
        self.image = m.imageActor()
        
        self.createViewer()
        
        self.setVolume(self.volume)
        self.setImage(self.image)
        
        self.addImage()
        self.addVolume()
        
        #self.viewer.scene.renderer.render()
        #self.viewer.scene.background = (1,1,1)
        #self.viewer.scene.renderer.reset_camera()
    
    def createImageViewer(self):
        """"A lazy way to set up a viewer with the default spine image"""
        m = myActors()
        self.image = m.imageActor()
        self.setImage(self.image)
        self.createViewer()
        self.addImage()
    
    def rotate(self, angle=90):
        """A little function to test out rendering capabilities"""
        #Actor must already be in the scene
        if self.viewer == []:
            print "viewer and viewer must be set before calling this method"
            return
        if self.volume == []:
            print "volume must be set before calling this method"
            return
        #self.viewer.scene.renderer.reset_camera()
        self.viewer.scene.camera.parallel_scale=100
        
        for i in range(0,angle):
            #The volume is rotated, not the camera
            self.volume.rotate_y(1)
            self.viewer.scene.render()
            
    def generateProjection(self,fn='output.png'):
        """Write an image ofthe scene to file"""
        #TODO: !!Allow the user to specify a transform to be applied to the volume!!
        self.viewer.open()
        w2if = tvtk.WindowToImageFilter()
        w2if.input = self.viewer.scene.render_window
        ex = tvtk.PNGWriter(input=w2if.output,file_name=fn)
        w2if.update()
        ex.write()
        
    def setViewer(self, viewer):
        self.viewer = viewer
    
    def setVolume(self,vol):
        self.volume = vol
        
    def setImage(self,image):
        self.image = image        
        
    def setStyle(self,style):
        """This function will allow the style to be changed.  Styles
        ending with 'actor' will interact with the actor, those 
        ending with 'camera' wll interact with the camera.  For JointTrack
        we will want to interact with the actors.  There is also an image
        interactor style that does not allow rotation of the images. """
        if self.viewer == []:
            self.createViewer()
        self.style = style
        self.viewer.scene.interactor.interactor_style = style
        
class myReg():
    def __init__(self):
        """"
        This class contains the function to do a registration between two images
        
        I tried to only setup the metric, but had errors, so i setup the whole registration.
        """        
        
        self.movingImage = []
        self.fixedImage = []
        self.internalImage = []
        
        #Setup image related ITK stuff
        dim = 2
        inputImageType = itk.Image[itk.US, dim]
        internalImageType = itk.Image[itk.F,dim]
        
        
        #Setup VTK-ITK conversion functions
        self.itk_vtk_converter = itk.ImageToVTKImageFilter[internalImageType].New()
        self.vtk_itk_converter = itk.VTKImageToImageFilter[internalImageType].New()
        
##        #Create Type caster function (unsigned short [input] to float[internal])
##        fixedCaster = itk.CastImageFilter[inputImageType, internalImageType].New()
##        movingCaster = itk.CastImageFilter[inputImageType, internalImageType].New()
##        
##        #Read some default images in for testing, these can be changed using 'setMovingImage' and 'setFixeImage' 
##        #fixedFileName = 'C:/Lib/ITK/InsightToolkit-3.6.0/examples/Data/BrainProtonDensitySliceBorder20.png'
##        fixedFileName = 'C:/Lib/JointTrack/branches/bryan/output.png'
##        fixedReader = itk.ImageFileReader[inputImageType].New()
##        fixedReader.SetFileName( fixedFileName )
##        fixedReader.Update()
##        fixedCaster.SetInput( fixedReader.GetOutput() )
##        self.fixedImage = fixedCaster.GetOutput()
##        #itk.echo(fixedCaster)
##        
##        #movingFileName = 'C:/Lib/ITK/InsightToolkit-3.6.0/examples/Data/BrainProtonDensitySliceR10X13Y17.png'
##        movingFileName = 'C:/Lib/JointTrack/branches/bryan/output_x15y15.png'        
##        movingReader = itk.ImageFileReader[inputImageType].New()
##        movingReader.SetFileName( movingFileName )
##        movingReader.Update()
##        movingCaster.SetInput( movingReader.GetOutput() )
##        self.movingImage = movingCaster.GetOutput()   
##        #itk.echo(movingCaster)
        
        #
        #  Read the fixed and moving images using filenames 
        #  from the command line arguments
        #
        fixedFileName = 'C:/Lib/JointTrack/branches/bryan/output.png'
        movingFileName = 'C:/Lib/JointTrack/branches/bryan/output_r10.png' 
        
        fixedImageReader  = itk.ImageFileReader.IF2.New()
        movingImageReader = itk.ImageFileReader.IF2.New()

        fixedImageReader.SetFileName(  fixedFileName )
        movingImageReader.SetFileName( movingFileName )

        fixedImageReader.Update()
        movingImageReader.Update()

        self.fixedImage  = fixedImageReader.GetOutput()
        self.movingImage = movingImageReader.GetOutput()
        
        print "module imported"
        
        
    def doReg(self):
        """
        Do the actual registration stuff
        """
        
        #
        #  Instantiate the classes for the registration framework
        #
        
        self.registration = itk.ImageRegistrationMethod.IF2IF2.New()
        imageMetric  = itk.MeanSquaresImageToImageMetric.IF2IF2.New()
        transform    = itk.CenteredRigid2DTransform.D.New()
        optimizer    = itk.RegularStepGradientDescentOptimizer.New()
        interpolator = itk.LinearInterpolateImageFunction.IF2D.New()
        
        self.registration.SetOptimizer(    optimizer.GetPointer()    )
        self.registration.SetTransform(    transform.GetPointer()    )
        self.registration.SetInterpolator( interpolator.GetPointer() )
        self.registration.SetMetric(       imageMetric.GetPointer()  )
        self.registration.SetFixedImage(  self.fixedImage  )
        self.registration.SetMovingImage( self.movingImage )
        
        self.registration.SetFixedImageRegion( self.fixedImage.GetBufferedRegion() )

        transform.SetIdentity()
        initialParameters = transform.GetParameters()

        self.registration.SetInitialTransformParameters( initialParameters )
        
        #Set the image center to the center of the image
        fixedImageCenterX = self.fixedImage.GetBufferedRegion().GetSize().GetElement(0)/2
        fixedImageCenterY = self.fixedImage.GetBufferedRegion().GetSize().GetElement(1)/2
        self.fixedImage.SetOrigin((fixedImageCenterX,fixedImageCenterY))
        print "Fixed image center: ", fixedImageCenterX, fixedImageCenterY

        movingImageCenterX = self.movingImage.GetBufferedRegion().GetSize().GetElement(0)/2
        movingImageCenterY = self.movingImage.GetBufferedRegion().GetSize().GetElement(1)/2
        self.movingImage.SetOrigin((movingImageCenterX,movingImageCenterY))
        print "Moving image center: ",movingImageCenterX, movingImageCenterY
        
        transform.SetCenter((fixedImageCenterX,fixedImageCenterY))
        
        #
        # Initial transform parameters 
        #
        transform.SetAngle( 0.0 );

        # center of the fixed image
        fixedSpacing = self.fixedImage.GetSpacing()
        fixedOrigin = self.fixedImage.GetOrigin()
        fixedSize = self.fixedImage.GetLargestPossibleRegion().GetSize()

        centerFixed = ( 256, 256 )
        self.fixedImage.SetOrigin(centerFixed)

        # center of the moving image 
        movingSpacing = self.movingImage.GetSpacing()
        movingOrigin = self.movingImage.GetOrigin()
        movingSize = self.movingImage.GetLargestPossibleRegion().GetSize()

        centerMoving = ( 256, 256 )
        self.movingImage.SetOrigin(centerMoving)

        # transform center
        center = transform.GetCenter()
        center.SetElement( 0, centerFixed[0] )
        center.SetElement( 1, centerFixed[1] )

        # transform translation
        #translation = transform.GetTranslation()
        #translation.SetElement( 0, centerMoving[0] - centerFixed[0] )
        #translation.SetElement( 1, centerMoving[1] - centerFixed[1] )

        initialParameters = transform.GetParameters()
        
        print "Image Parameters: "
        print "     Moving Spacing: ", movingSpacing.GetElement(0), movingSpacing.GetElement(1)
        print "     Fixed Spacing: ", fixedSpacing.GetElement(0), fixedSpacing.GetElement(1)
        print "     Moving Size: ", movingSize.GetElement(0), movingSize.GetElement(1)
        print "     Fixed Size: ", fixedSize.GetElement(0), fixedSize.GetElement(1)
        
        print "Initial Parameters: "
        print "     Angle: %f" % (initialParameters.GetElement(0), )
        print "     Center: %f, %f" % ( initialParameters.GetElement(1), initialParameters.GetElement(2) )
        print "     Translation: %f, %f" % (initialParameters.GetElement(3), initialParameters.GetElement(4))
        
        print "Number of pixels counted: ", imageMetric.GetNumberOfPixelsCounted()
        self.registration.SetInitialTransformParameters( initialParameters )
        
        #print "Metric Value: ", imageMetric.GetValue((0,0,0,0,0))
        
        #print "Transform: ", itk.echo(transform)
        #
        # Iteration Observer
        #
        def iterationUpdate():
            currentParameter = transform.GetParameters()
            transform.SetParameters((10,256,256,0,0))
            print "M: %f   P: %f %f %f %f %f" % ( optimizer.GetValue(),
                                currentParameter.GetElement(0),
                                currentParameter.GetElement(1),
                                currentParameter.GetElement(2),
                                currentParameter.GetElement(3),
                                currentParameter.GetElement(4) )
         

        iterationCommand = itk.PyCommand.New()
        iterationCommand.SetCommandCallable( iterationUpdate )
        optimizer.AddObserver( itk.IterationEvent(), iterationCommand.GetPointer() )

        #
        #  Define optimizer parameters
        #
        optimizer.SetMaximumStepLength(  4.00 )
        optimizer.SetMinimumStepLength(  0.01 )
        optimizer.SetNumberOfIterations( 200  )


        
        #
        #  Start the registration process
        #

##        currentParameter = transform.GetParameters()                        
##        print "DEBUG metric initial value: ", optimizer.GetValue()
##        print "DEBUG transform parameters: ", currentParameter.GetElement(0), currentParameter.GetElement(1)
        #print "Fixed Image ..."
        #itk.echo(self.fixedImage)
        #print "Moving Image ..."
        #itk.echo(self.movingImage)
        #print "Transform ..."
        #itk.echo(transform)
        
        print "Starting registration"
        self.registration.StartRegistration() 


        #
        # Get the final parameters of the transformation
        #
        finalParameters = self.registration.GetLastTransformParameters()

        print "Final Registration Parameters "
        print "Translation X =  %f" % (finalParameters.GetElement(0),)
        print "Translation Y =  %f" % (finalParameters.GetElement(1),)
        transform.SetParameters((10,256,256,0,0))
        #print "Metric Value: ", optimizer.GetValue()
    
    def getOutputImage(self):
        """
        resample and write output image to file
        """
        # Now, we use the final transform for resampling the moving image.
        resampler = itk.ResampleImageFilter.IF2IF2.New()

        resampler.SetTransform( self.registration.GetTransform() )
        resampler.SetInput( self.movingImage )

        region = self.fixedImage.GetLargestPossibleRegion()

        resampler.SetSize( region.GetSize() )
        resampler.SetOutputSpacing( self.fixedImage.GetSpacing() )
        resampler.SetOutputDirection( self.fixedImage.GetDirection() )
        resampler.SetOutputOrigin(  self.fixedImage.GetOrigin() )
        resampler.SetDefaultPixelValue( 100 )

        #
        # Cast for output
        #
        outputCast = itk.RescaleIntensityImageFilter.IF2IUS2.New()
        
        outputCast.SetInput( resampler.GetOutput() )
        outputCast.SetOutputMinimum( 0 )
        outputCast.SetOutputMaximum( 65535 )

        writer = itk.ImageFileWriter.IUS2.New()

        writer.SetFileName( "C:/Lib/JointTrack/branches/bryan/reg_output.png" )
        writer.SetInput( outputCast.GetOutput() )
        writer.Update()
        
    def setMovingImage(self,img):
        """Set the moving image for registration"""
        self.movingImage=img
        self.metric.SetMovingImage( self.movingImage )
        
    def setFixedImage(self,img):
        """Set the moving image for registration"""
        self.fixedImage=img
        self.metric.SetFixedImage (self.fixedImage ) 
        